#!/usr/bin/env python3
import os
from collections import defaultdict
import tabulate

def count_lines_in_file(file_path):
    with open(file_path, "r", encoding="utf-8", errors="ignore") as file:
        return sum(1 for line in file)


def count_code_files(directory, extensions, depth=1):
    total_lines = 0
    lines_per_folder = defaultdict(int)
    lines_per_folder_per_extension = defaultdict(lambda: defaultdict(int))  # 新增字典来统计每个文件夹下不同后缀的文件行数

    # 获取当前一层的子文件夹路径
    subdirectories = [
        os.path.join(directory, d)
        for d in os.listdir(directory)
        if os.path.isdir(os.path.join(directory, d))
    ]

    header = "undefine"

    for root, _, files in os.walk(directory):
        current_depth = root.count(os.sep) - directory.count(os.sep)
        if current_depth == depth:
            header = os.path.basename(root)

        for file in files:
            if any(file.endswith(ext) for ext in extensions):
                file_path = os.path.join(root, file)
                try:
                    lines = count_lines_in_file(file_path)
                    total_lines += lines
                    lines_per_folder[header] += lines
                    lines_per_folder_per_extension[header][file.split('.')[-1]] += lines
                except Exception as e:
                    print(f"Error reading {file_path}: {e}")

    return total_lines, lines_per_folder, lines_per_folder_per_extension


if __name__ == "__main__":
    project_directory = "./package"
    depth = 1
    code_extensions = [".py", ".cpp", ".txt", ".xml"]

    total_lines, lines_per_folder, lines_per_folder_per_extension = count_code_files(
        project_directory, code_extensions, depth
    )

    print("Total lines of code:", total_lines)
    headers = ["Folder"] + code_extensions + ["Total"]
    table_data = []

    for folder, extensions in lines_per_folder_per_extension.items():
        row = [folder]
        folder_total = 0
        for ext in code_extensions:
            lines = extensions.get(ext[1:], 0)
            row.append(lines)
            folder_total += lines
        row.append(folder_total)
        table_data.append(row)

    # 对 table_data 按照每个文件夹的总行数降序排序
    table_data.sort(key=lambda x: x[-1], reverse=True)

    print(tabulate.tabulate(table_data, headers=headers, tablefmt="grid"))